# -------------------------1.导入头文件---------------------------------
import warnings
warnings.filterwarnings("ignore")
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as T
from torchvision import models

from backbone.resnet import resnet50
from backbone.loss import Yolo_Loss
from utils.Yolo_Dataset import Yolo_Dataset

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
batch_size = 4      # 根据自己的电脑设定
epochs = 1
lr = 0.01
file_root = './VOCdevkit/VOC2007/JPEGImages'   # 需要根据的实际路径修改


# -------------3. 创建模型并继承预训练参数------------------
# 加载预训练模型
# 接下来就是让自己的模型去继承权重参数
model = resnet50()  # 创建模型实例
# model.load_state_dict(torch.load('D:\Pycharm\yolov1\save_weights\yolo.pth'))  # 加载权重
# print('预训练模型加载完成')
model.to(device)  # 转移到设备

# -------------------4. 将模型等放入GPU中--------------
loss = Yolo_Loss()
optimizer = torch.optim.SGD(model.parameters(),lr=lr,momentum=0.9,weight_decay=5e-4)
model.to(device)
loss.to(device)
# 5. 加载数据
train_dataset = Yolo_Dataset(root=file_root,list_file='./utils/voctrain.txt',train=True,transforms = [T.ToTensor()])
train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True,drop_last=True)
test_dataset = Yolo_Dataset(root=file_root,list_file='./utils/voctest.txt',train=False,transforms = [T.ToTensor()])
test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=True,drop_last=True)

# -------------------------5. 训练 ---------------------------
# 打印一些基本的信息
print('starting train the model')
print('the train_dataset has %d images' % len(train_dataset))
print('the batch_size is ',batch_size)
# 定义一个最佳损失值
best_test_loss = 0
# 开始训练
for e in range(epochs):
    model.train()
    # 调整学习率
    if e == 20:
        print('change the lr')
        optimizer.param_groups[0]['lr'] /= 10
    if e == 35:
        print('change the lr')
        optimizer.param_groups[0]['lr'] /= 10
    # 进度条显示
    tqdm_tarin = tqdm(train_loader)
    # 定义损失变量
    total_loss = 0.
    for i,(images,target) in enumerate(tqdm_tarin):
        # 将变量放入设备中
        images,target = images.to(device),target.to(device)
        # 训练--损失等
        pred = model(images)
        loss_value = loss(pred,target)
        total_loss += loss_value.item()
        optimizer.zero_grad()
        loss_value.backward()
        optimizer.step()
        # 打印一下损失值
        if (i+1) % 5 == 0:
            tqdm_tarin.desc = 'train epoch[{}/{}] loss:{:.6f}'.format(e+1,epochs,total_loss/(i+1))
    # 启用验证模式
    model.eval()
    validation_loss = 0.0
    tqdm_test = tqdm(test_loader)
    for i, (images, target) in enumerate(tqdm_test):
        images, target = images.cuda(), target.cuda()
        pred = model(images)
        loss_value = loss(pred, target)
        validation_loss += loss_value.item()
    validation_loss /= len(test_loader)
    # 显示验证集的损失值
    print('In the test step,the average loss is %.6f' % validation_loss)
    # --------------------- 6.保存模型---------------------------------
    if best_test_loss > validation_loss:
        best_test_loss = validation_loss
        print('get best test loss %.5f' % best_test_loss)
        torch.save(model.state_dict(), './save/best.pth')
    # 记得最后保存参数
    torch.save(model.state_dict(), './save/yolo1.pth')

